package dict

import (
	"sync"
)

// BiDict is goroutine safe bidirectional map
type BiDict[K comparable, V comparable] struct {
	mutex sync.RWMutex
	k2v   map[K]V
	v2k   map[V]K
}

func NewBiDict[K comparable, V comparable]() *BiDict[K, V] {
	p := new(BiDict[K, V])
	p.k2v = make(map[K]V)
	p.v2k = make(map[V]K)
	return p
}

func (p *BiDict[K, V]) Set(key K, val V) {
	p.mutex.Lock()
	defer p.mutex.Unlock()
	p.k2v[key] = val
	p.v2k[val] = key
}

func (p *BiDict[K, V]) GetByKey(key K) (val V, found bool) {
	p.mutex.RLock()
	defer p.mutex.RUnlock()
	val, found = p.k2v[key]
	return
}

func (p *BiDict[K, V]) GetByVal(val V) (key K, found bool) {
	p.mutex.RLock()
	defer p.mutex.RUnlock()
	key, found = p.v2k[val]
	return
}

func (p *BiDict[K, V]) DeleteByKey(key K) {
	p.mutex.Lock()
	defer p.mutex.Unlock()
	if val, found := p.k2v[key]; found {
		delete(p.v2k, val)
	}
	delete(p.k2v, key)
}

func (p *BiDict[K, V]) DeleteByVal(val V) {
	p.mutex.Lock()
	defer p.mutex.Unlock()
	if key, found := p.v2k[val]; found {
		delete(p.k2v, key)
	}
	delete(p.v2k, val)
}

func (p *BiDict[K, V]) Clear() {
	p.mutex.Lock()
	defer p.mutex.Unlock()
	p.k2v = make(map[K]V)
	p.v2k = make(map[V]K)
}

func (p *BiDict[K, V]) Len() int {
	p.mutex.RLock()
	defer p.mutex.RUnlock()
	return len(p.k2v)
}

func (p *BiDict[K, V]) Clone() (k2v map[K]V, v2k map[V]K) {
	p.mutex.RLock()
	defer p.mutex.RUnlock()
	k2v = make(map[K]V, len(p.k2v))
	for key, val := range p.k2v {
		k2v[key] = val
	}

	v2k = make(map[V]K, len(p.v2k))
	for val, key := range p.v2k {
		v2k[val] = key
	}
	return
}

// Range iter element of p with f
// if f return false, then break the range
// note: do not call any p's function inside f, otherwise will occur deadlock
func (p *BiDict[K, V]) Range(f func(K, V) bool) {
	p.mutex.RLock()
	defer p.mutex.RUnlock()
	for key, val := range p.k2v {
		if f(key, val) == false {
			break
		}
	}
}
